#include "hook.h"
#include <dlfcn.h>
#include "fiber.h"
#include "iomanager.h"
#include "fd_manager.h"
#include "util.h"
#include "log.h"
#include "config.h"
#include "scheduler.h"

yuechuan::Logger::ptr g_logger = YUECHUAN_LOG_NAME("root");
namespace yuechuan {
    static thread_local bool t_hook_enable = false;
    static uint64_t s_connect_timeout = -1;
    /// socket 中的connect超时变量
    static yuechuan::ConfigVar<int>::ptr g_tcp_connect_timeout = yuechuan::Config::Lookup("tcp.connect.timeout", 5000,
                                                                                          "tcp connect timeout");
#define HOOK_FUN(XX)\
    XX(sleep)\
    XX(usleep)\
    XX(nanosleep)\
    XX(socket)\
    XX(connect)\
    XX(accept)\
    XX(read)\
    XX(readv)\
    XX(recv)\
    XX(recvfrom)\
    XX(recvmsg)\
    XX(write)\
    XX(writev)\
    XX(send)\
    XX(sendto)\
    XX(sendmsg)\
    XX(close)\
    XX(fcntl)\
    XX(ioctl)\
    XX(getsockopt)\
    XX(setsockopt)

    void hook_init() {
        static bool is_inited = false;
        if (is_inited) {
            return;
        }
#define XX(name) name ## _f = (name ## _fun)dlsym(RTLD_NEXT, #name);
        HOOK_FUN(XX);
#undef XX
    }

    struct HookIniter {
        HookIniter() {
            hook_init();
            s_connect_timeout = g_tcp_connect_timeout->getValue();
            g_tcp_connect_timeout->addListener([](const int &oldvalue, const int &newvalue) {
                YUECHUAN_LOG_INFO(g_logger) << "tcp connect timeout changed from " << oldvalue << "to " << newvalue;
                s_connect_timeout = newvalue;
            });

        }
    };

    static HookIniter s_hook_initer;


    bool is_hook_enable() {
        return t_hook_enable;
    }

    void set_hook_enable(bool flag) {
        t_hook_enable = flag;
    }
}
/**
 * @brief 用于创建条件定时器的条件
 */
struct timer_info {
    int canceled = 0;
};

/*
 *大部分操作都是和文件的IO操作有关的，可以将其分离处一个函数用于处理文件句柄要执行的函数
 * 对于socket类型的文件要特殊管理执行实现的非阻塞方法
 */
template<typename OriginFun, typename...Args>
static ssize_t do_io(int fd, OriginFun fun, const char *hook_fun_name,
                     uint32_t event, int timeout_so, Args &&... args) {
    if (!yuechuan::t_hook_enable) {
        return fun(fd, std::forward<Args>(args)...);    // 首先要判断执行io操作的文件是否被hook住，没有直接执行其原始的方法
    }
    yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);  // 获取当前文件的句柄上下文管理的对象
    if (!ctx) {// 不存在说明该文件句柄不是socket，那么就要调用其原先的方法
        return fun(fd, std::forward<Args>(args)...);
    }
    if (ctx->isClose()) {                                    // 当文件状态为close时，需要改变错误代码为BADFILE，返回-1表示错误
        errno = EBADF;
        return -1;
    }
    if (!ctx->isSocket() || ctx->getUserNonblock()) {        // 当文件类型不是socket类型，或者用户自行设置其状态为非阻塞，直接执行其原始的回调方法
        return fun(fd, std::forward<Args>(args)...);
    }
    uint64_t to = ctx->getTimeout(timeout_so);
    std::shared_ptr<timer_info> tinfo(new timer_info);
    ssize_t n;
    while (true) {
        n = fun(fd, std::forward<Args>(args)...);
        while (n == -1 && errno == EINTR) {                    // 状态是中断，则进行重试
            n = fun(fd, std::forward<Args>(args)...);
        }
        if (n == -1 && errno == EAGAIN) {                        // 再次错误，说明其正在阻塞状态，要再次执行io操作
//            YUECHUAN_LOG_INFO(g_logger) << "do_io " << hook_fun_name;
            yuechuan::IOManager *iom = yuechuan::IOManager::GetThis();

            yuechuan::Timer::ptr timer;
            std::weak_ptr<timer_info> winfo(tinfo);

            if (to != (uint64_t) -1) {
                timer = iom->addConditionTimer(to, [winfo, fd, iom, event]() { //设置超时事件
                    auto t = winfo.lock();
                    if (!t || t->canceled) {
                        // 当lock的返回值为false时，或者手动对canceled进行赋值，则表明该文件句柄的函数放弃执行
                        return;
                    }
                    t->canceled = ETIMEDOUT;
                    iom->cancelEvent(fd, (yuechuan::IOManager::Event) (event));          // 取消事件，并强制唤醒
                }, winfo);
            }
            int rt = iom->addEvent(fd, (yuechuan::IOManager::Event) (event));              // 设置定时器事件
            if (rt) {
                // 当addEvent返回值非0时，表示添加事件错误
                YUECHUAN_LOG_ERROR(g_logger) << hook_fun_name << "addEvent( " << fd << ". " << event << " )";

                if (timer) {
                    // 前面已经添加了事件的定时器，当添加事件失败后需要将该定时器取消
                    timer->cancel();
                }
                return -1;
            } else {
                yuechuan::Fiber::YieldToHold();
                if (timer) {
                    timer->cancel();
                }
                if (tinfo->canceled) {
                    errno = tinfo->canceled;
                    return -1;
                }
                continue;
            }
        }
        return n;
    }

}

extern "C" {
// name_fun name_f = nullptr;
#define XX(name) name ## _fun name ## _f = nullptr;
HOOK_FUN(XX) ;
#undef XX
unsigned int sleep(unsigned int seconds) {
    if (!yuechuan::t_hook_enable) {
        return sleep_f(seconds);
    }
    yuechuan::Fiber::ptr fiber = yuechuan::Fiber::GetThis();
    yuechuan::IOManager *iom = yuechuan::IOManager::GetThis();
    iom->addTimer(seconds * 1000, std::bind(
            (void (yuechuan::Scheduler::*)(yuechuan::Fiber::ptr, int thread)) &yuechuan::IOManager::schedule,
            iom, fiber, -1));
    yuechuan::Fiber::YieldToHold();
    return 0;

}
int usleep(useconds_t usec) {
    if (!yuechuan::t_hook_enable) {
        return usleep_f(usec);
    }
    yuechuan::Fiber::ptr fiber = yuechuan::Fiber::GetThis();
    yuechuan::IOManager *iom = yuechuan::IOManager::GetThis();

    iom->addTimer(usec / 1000, std::bind(
            (void (yuechuan::Scheduler::*)(yuechuan::Fiber::ptr, int thread)) &yuechuan::IOManager::schedule,
            iom, fiber, -1));
    yuechuan::Fiber::YieldToHold();

    return 0;
}

int nanosleep(const struct timespec *req, struct timespec *rem) {
    if (!yuechuan::t_hook_enable) {
        return nanosleep_f(req, rem);
    }
    int timeout_ms = req->tv_sec * 1000 + req->tv_nsec / 1000 / 1000;
    yuechuan::Fiber::ptr fiber = yuechuan::Fiber::GetThis();
    yuechuan::IOManager *iom = yuechuan::IOManager::GetThis();
    iom->addTimer(timeout_ms, [iom, fiber]() {
        iom->schedule(fiber);
    });
    yuechuan::Fiber::YieldToHold();
    return 0;

}

int socket(int domain, int type, int protocol) {
    if (!yuechuan::t_hook_enable) {
        return socket_f(domain, type, protocol);
    }
    int fd = socket_f(domain, type, protocol);
    if (fd == -1) {
        return fd;
    }
    yuechuan::FdMgr::GetInstance()->get(fd, true);
    return fd;

}

int connect_with_timeout(int fd, const struct sockaddr *addr, socklen_t addrlen, uint64_t timeout_ms) {
    if (!yuechuan::t_hook_enable) {
        return connect_f(fd, addr, addrlen);
    }
    yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);
    if (!ctx || ctx->isClose()) {
        errno = EBADF;
        return -1;
    }
    if (!ctx->isSocket()) {
        return connect_f(fd, addr, addrlen);
    }
    if (ctx->getUserNonblock()) {
        return connect_f(fd, addr, addrlen);
    }
    int n = connect_f(fd, addr, addrlen);
    if (n == 0) {
        return 0;
    } else if (n != -1 || errno != EINPROGRESS) {
        return n;
    }
    yuechuan::IOManager *iom = yuechuan::IOManager::GetThis();
    yuechuan::Timer::ptr timer;
    std::shared_ptr<timer_info> tinfo(new timer_info);
    std::weak_ptr<timer_info> winfo(tinfo);

    if (timeout_ms != (uint64_t) -1) {
        timer = iom->addConditionTimer(timeout_ms, [winfo, fd, iom]() {
            auto t = winfo.lock();
            if (!t || t->canceled) {
                return;
            }
            t->canceled = ETIMEDOUT;
            iom->cancelEvent(fd, yuechuan::IOManager::WRITE);
        }, winfo);
    }
    int rt = iom->addEvent(fd, yuechuan::IOManager::WRITE);
    if (rt == 0) {
        yuechuan::Fiber::YieldToHold();
        if (timer) {
            timer->cancel();
        }
        if (tinfo->canceled) {
            errno = tinfo->canceled;
            return -1;
        }
    } else {
        if (timer) {
            timer->cancel();
        }
        YUECHUAN_LOG_ERROR(g_logger) << " connect addEvent {" << fd << ", WRITE} ERROR";
    }
    int error = 0;
    socklen_t len = sizeof(int);
    if (-1 == getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len)) {
        return -1;
    }
    if (!error) {
        return 0;
    } else {
        errno = error;
        return -1;
    }
}

int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
    return connect_with_timeout(sockfd, addr, addrlen, yuechuan::s_connect_timeout);
//         return connect_f(sockfd, addr, addrlen);

}


int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
    int fd = do_io(sockfd, accept_f, "accept", yuechuan::IOManager::READ, SO_RCVTIMEO, addr, addrlen);
    if (fd >= 0) {
        yuechuan::FdMgr::GetInstance()->get(fd, true);
    }
    return fd;
}

ssize_t read(int fd, void *buf, size_t count) {
    return do_io(fd, read_f, "read", yuechuan::IOManager::READ, SO_RCVTIMEO, buf, count);
}

ssize_t readv(int fd, const struct iovec *iov, int iovcnt) {
    return do_io(fd, readv_f, "readv", yuechuan::IOManager::READ, SO_RCVTIMEO, iov, iovcnt);
}


ssize_t recv(int sockfd, void *buf, size_t len, int flags) {
    return do_io(sockfd, recv_f, "recv", yuechuan::IOManager::READ, SO_RCVTIMEO, buf, len, flags);
}

ssize_t recvfrom(int sockfd, void *buf, size_t len, int flags, struct sockaddr *src_addr, socklen_t *addrlen) {
    return do_io(sockfd, recvfrom_f, "recvfrom", yuechuan::IOManager::READ, SO_RCVTIMEO, buf, len, flags, src_addr,
                 addrlen);
}

ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags) {
    return do_io(sockfd, recvmsg_f, "recvmsg", yuechuan::IOManager::READ, SO_RCVTIMEO, msg, flags);
}

ssize_t write(int fd, const void *buf, size_t count) {
    return do_io(fd, write_f, "write", yuechuan::IOManager::WRITE, SO_SNDTIMEO, buf, count);
}


ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
    return do_io(fd, writev_f, "writev", yuechuan::IOManager::WRITE, SO_SNDTIMEO, iov, iovcnt);
}


ssize_t send(int sockfd, const void *buf, size_t len, int flags) {
    return do_io(sockfd, send_f, "send", yuechuan::IOManager::WRITE, SO_SNDTIMEO, buf, len, flags);
}


ssize_t
sendto(int sockfd, const void *buf, size_t len, int flags, const struct sockaddr *dest_addr, socklen_t addrlen) {
    return do_io(sockfd, sendto_f, "sendto", yuechuan::IOManager::WRITE, SO_SNDTIMEO, buf, len, flags, dest_addr,
                 addrlen);
}


ssize_t sendmsg(int sockfd, const struct msghdr *msg, int flags) {
    return do_io(sockfd, sendmsg_f, "sendmsg", yuechuan::IOManager::WRITE, SO_SNDTIMEO, msg, flags);
}
int close(int fd) {
    if (!yuechuan::t_hook_enable) {
        return close_f(fd);
    }
    yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);
    if (ctx) {
        auto iom = yuechuan::IOManager::GetThis();
        if (iom) {
            iom->cancelAll(fd);
        }
        yuechuan::FdMgr::GetInstance()->del(fd);
    }
    return close_f(fd);

}

int fcntl(int fd, int cmd, ... /* arg */ ) {
    va_list va;
    va_start(va, cmd);
    switch (cmd) {
        case F_SETFL: {
            int arg = va_arg(va, int);
            va_end(va);
            yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);
            if (!ctx || ctx->isClose() || !ctx->isSocket()) {
                return fcntl_f(fd, cmd, arg);
            }
            ctx->setUserNonblock(arg & O_NONBLOCK);
            if (ctx->getSysNonblock()) {
                arg |= O_NONBLOCK;
            } else {
                arg &= ~O_NONBLOCK;
            }
            return fcntl_f(fd, cmd, arg);
        }
            break;
        case F_GETFL: {
            va_end(va);
            int arg = fcntl_f(fd, cmd);
            yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);
            if (!ctx || ctx->isClose() || !ctx->isSocket()) {
                return arg;
            }
            if (ctx->getUserNonblock()) {
                return arg | O_NONBLOCK;
            } else {
                return arg & ~O_NONBLOCK;
            }
        }
            break;
        case F_DUPFD:
        case F_DUPFD_CLOEXEC:
        case F_SETFD:
        case F_SETOWN:
        case F_SETSIG:
        case F_SETLEASE:
        case F_NOTIFY:
#ifdef F_SETPIPE_SZ
        case F_SETPIPE_SZ:
#endif
        {
            int arg = va_arg(va, int);
            va_end(va);
            return fcntl_f(fd, cmd, arg);
        }
            break;
        case F_GETFD:
        case F_GETOWN:
        case F_GETSIG:
        case F_GETLEASE:
#ifdef F_GETPIPE_SZ
        case F_GETPIPE_SZ:
#endif
        {
            va_end(va);
            return fcntl_f(fd, cmd);
        }
            break;
        case F_SETLK:
        case F_SETLKW:
        case F_GETLK: {
            struct flock *arg = va_arg(va, struct flock*);
            va_end(va);
            return fcntl_f(fd, cmd, arg);
        }
            break;
        case F_GETOWN_EX:
        case F_SETOWN_EX: {
            struct f_owner_exlock *arg = va_arg(va, struct f_owner_exlock*);
            va_end(va);
            return fcntl_f(fd, cmd, arg);
        }
            break;
        default:
            va_end(va);
            return fcntl_f(fd, cmd);
    }
}

int ioctl(int fd, unsigned long request, ...) {
    va_list va;
    va_start(va, request);
    void *arg = va_arg(va, void*);
    va_end(va);
    if (FIONBIO == request) {
        bool user_nonblock = !!*(int *) arg;
        yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(fd);
        if (!ctx || ctx->isClose() || !ctx->isSocket()) {
            return ioctl_f(fd, request, arg);
        }
        ctx->setUserNonblock(user_nonblock);
    }
    return ioctl_f(fd, request, arg);

}


int getsockopt(int sockfd, int level, int optname, void *optval, socklen_t *optlen) {
    return getsockopt_f(sockfd, level, optname, optval, optlen);
}
int setsockopt(int sockfd, int level, int optname, const void *optval, socklen_t optlen) {
    if (!yuechuan::t_hook_enable) {
        return setsockopt_f(sockfd, level, optname, optval, optlen);
    }
    if (level == SOL_SOCKET) {
        if (optname == SO_RCVTIMEO || optname == SO_SNDTIMEO) {
            yuechuan::FdCtx::ptr ctx = yuechuan::FdMgr::GetInstance()->get(sockfd);
            if (ctx) {
                const timeval *tv = (const timeval *) optval;
                ctx->setTimeout(optname, tv->tv_sec * 1000 + tv->tv_usec / 1000);
            }
        }
    }
    return setsockopt_f(sockfd, level, optname, optval, optlen);
}

}
